import cv2
import numpy as np
import time
import threading
import serial
from serial.tools.list_ports import comports
import struct

# ----------------------
# 配置参数
# ----------------------
SERAIL = "/dev/ttyAMA0"  # 注意修正串口号
BAUDTRATE = 115200
RESOLUTION = (640, 480)  # 帧宽，帧高
FRAMERATE = 15           # 帧率
SHAPE_DETECTION_THRESH = 0.7  # 置信度阈值
MIN_CONTOUR_AREA = 100   # 最小轮廓面积
SEND_INTERVAL = 1.0      # 数据发送间隔
DEBUG_MODE = True        # 是否显示调试窗口

# ----------------------
# 映射字典
# ----------------------
dic_object1 = {
    0: 'bus',
    1: 'car',
    2: 'cat',
    3: 'dog',
    4: 'pig',
    5: 'truck',
    6: 'blank'
}
dic_object2 = {v: k+1 for k, v in dic_object1.items()}  # {'bus':1, ...}
dic_shape = {
    'circle': 1,
    'rectangle': 2,
}

# ----------------------
# UART 类（修复 is_open）
# ----------------------
class UART:
    def __init__(self, port, baudrate):
        self.port = port
        self.baudrate = baudrate
        self.connected = False
        try:
            self.ser = serial.Serial(port=self.port, baudrate=self.baudrate, timeout=1)
            self.connected = self.ser.is_open
            if self.connected:
                print(f"打开串口成功: {self.ser.name}")
        except serial.SerialException as e:
            print(f"打开串口失败: {e}")

    def uart_send_packet(self, data_object, data_shape):
        header = 0xA5
        footer = 0x5A
        fmt = '<BBBB'  # 小端序，4个 unsigned byte
        try:
            data_packet = struct.pack(fmt, header, data_object, data_shape, footer)
            self.ser.write(data_packet)
            print(f"UART 已发送数据包: {list(data_packet)}")
        except Exception as e:
            print(f"发送数据包失败: {e}")

    def async_send_serial(self, data_object, data_shape):
        thread = threading.Thread(
            target=self.uart_send_packet,
            args=(data_object, data_shape),
            daemon=True
        )
        thread.start()

# ----------------------
# 边框绘制函数
# ----------------------
def plot_one_box(xyxy, img, color=(0, 255, 0), label=None, line_thickness=3):
    c1, c2 = (int(xyxy[0]), int(xyxy[1])), (int(xyxy[2]), int(xyxy[3]))
    cv2.rectangle(img, c1, c2, color, thickness=line_thickness, lineType=cv2.LINE_AA)
    if label:
        tf = max(line_thickness - 1, 1)
        t_size = cv2.getTextSize(label, 0, fontScale=line_thickness / 3, thickness=tf)[0]
        c2 = c1[0] + t_size[0], c1[1] - t_size[1] - 3
        cv2.rectangle(img, c1, c2, color, -1, cv2.LINE_AA)
        cv2.putText(img, label, (c1[0], c1[1] - 2), 0, line_thickness / 3, [255, 255, 255], thickness=tf, lineType=cv2.LINE_AA)

# ----------------------
# 后处理函数（NMS）
# ----------------------
def post_process_opencv(outputs, model_h, model_w, img_h, img_w, thred_nms=0.4, thred_cond=0.5):
    conf = outputs[:, 4]
    c_x = outputs[:, 0] / model_w * img_w
    c_y = outputs[:, 1] / model_h * img_h
    w = outputs[:, 2] / model_w * img_w
    h = outputs[:, 3] / model_h * img_h
    p_cls = outputs[:, 5:]
    cls_id = np.argmax(p_cls, axis=1)
    p_x1 = c_x - w / 2
    p_y1 = c_y - h / 2
    p_x2 = c_x + w / 2
    p_y2 = c_y + h / 2
    areas = np.stack([p_x1, p_y1, p_x2, p_y2], axis=-1)
    ids = cv2.dnn.NMSBoxes(areas.tolist(), conf.tolist(), thred_cond, thred_nms)
    if len(ids) > 0:
        ids = ids.flatten()
        return areas[ids], conf[ids], cls_id[ids]
    return np.array([]), np.array([]), np.array([])

# ----------------------
# 推理函数
# ----------------------
def infer_image(net, img, model_h, model_w, thred_nms=0.4, thred_cond=0.5):
    if img is None or img.size == 0:
        return np.array([]), np.array([]), np.array([])
    img0 = img.copy()
    img = cv2.resize(img, (model_w, model_h))
    blob = cv2.dnn.blobFromImage(img, scalefactor=1/255.0, size=(model_w, model_h), swapRB=True, crop=False)
    net.setInput(blob)
    try:
        outs = net.forward()[0]
    except Exception as e:
        print(f"[ERROR] 推理失败: {e}")
        return np.array([]), np.array([]), np.array([])
    return post_process_opencv(outs, model_h, model_w, img0.shape[0], img0.shape[1], thred_nms, thred_cond)

# ----------------------
# 形状检测
# ----------------------
def detect_shape_simple(roi):
    try:
        gray = cv2.cvtColor(roi, cv2.COLOR_BGR2GRAY)
        _, thresh = cv2.threshold(gray, 127, 255, cv2.THRESH_BINARY_INV)
        contours, _ = cv2.findContours(thresh, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
        if not contours:
            return "unknown"
        cnt = max(contours, key=cv2.contourArea)
        if cv2.contourArea(cnt) < MIN_CONTOUR_AREA:
            return "unknown"
        peri = cv2.arcLength(cnt, True)
        approx = cv2.approxPolyDP(cnt, 0.04 * peri, True)
        if len(approx) == 4:
            return "rectangle"
        elif len(approx) > 8:
            return "circle"
    except Exception as e:
        print(f"[ERROR] 形状识别失败: {e}")
    return "unknown"

# ----------------------
# 主程序入口
# ----------------------
def main(uart_id, baudrate):
    my_uart = UART(uart_id, baudrate)

    # 加载模型
    model_path = '/home/txw/picture_find/picture_find.onnx'
    try:
        net = cv2.dnn.readNet(model_path)
        print("模型加载成功")
    except Exception as e:
        print(f" 模型加载失败: {e}")
        return

    # 打开摄像头
    cap = cv2.VideoCapture(0)
    if not cap.isOpened():
        print(" 摄像头无法打开")
        return

    cap.set(cv2.CAP_PROP_FRAME_WIDTH, RESOLUTION[0])
    cap.set(cv2.CAP_PROP_FRAME_HEIGHT, RESOLUTION[1])
    cap.set(cv2.CAP_PROP_FPS, FRAMERATE)

    last_send_time = time.time()

    try:
        while True:
            ret, frame = cap.read()
            if not ret:
                print("视频流中断，尝试重连...")
                cap.release()
                time.sleep(1)
                cap = cv2.VideoCapture(0)
                continue

            # 跳帧处理（每两帧处理一次）
            if frame_counter % 2 != 0:
                frame_counter += 1
                continue
            frame_counter = 0

            # 推理
            det_boxes, scores, ids = infer_image(net, frame, 640, 640)

            detections = []
            if len(ids) > 0:
                for box, score, id in zip(det_boxes, scores, ids):
                    if score < SHAPE_DETECTION_THRESH:
                        continue
                    label = dic_object1.get(int(id), "unknown")
                    x1, y1, x2, y2 = map(int, box)
                    x1, y1 = max(0, x1), max(0, y1)
                    x2, y2 = min(frame.shape[1]-1, x2), min(frame.shape[0]-1, y2)
                    roi = frame[y1:y2, x1:x2]
                    shape = detect_shape_simple(roi)
                    detections.append((label, float(score), box, shape))

                    if DEBUG_MODE:
                        plot_one_box(box, frame, label=f"{label} {score:.2f} {shape}", color=(0, 255, 0))

            # 发送最高置信度目标
            if detections and time.time() - last_send_time >= SEND_INTERVAL:
                detections.sort(key=lambda x: x[1], reverse=True)
                top_label, _, _, top_shape = detections[0]
                object_id = dic_object2.get(top_label, 0)
                shape_id = dic_shape.get(top_shape, 0)
                my_uart.async_send_serial(object_id, shape_id)
                last_send_time = time.time()

            if DEBUG_MODE:
                cv2.imshow('Frame', frame)
                if cv2.waitKey(1) == ord('q'):
                    break

    finally:
        cap.release()
        cv2.destroyAllWindows()

if __name__ == "__main__":
    main(SERAIL, BAUDTRATE)